package eventBusV3

import (
	"errors"
	"sync"
)

const (
	numWorkers    = 4
	taskQueueSize = 256
)

type Bus struct {
	subNode     sync.Map
	taskQueue   chan *asyncTask // 修改任务队列的类型
	shutdownSig chan struct{}
	taskPool    sync.Pool
}

func NewBus() *Bus { // 修改返回类型
	bus := &Bus{
		taskQueue:   make(chan *asyncTask, taskQueueSize), // 修改任务队列的类型
		shutdownSig: make(chan struct{}),
		taskPool: sync.Pool{
			New: func() interface{} {
				return &asyncTask{}
			},
		},
	}
	for i := 0; i < numWorkers; i++ {
		go bus.worker()
	}
	return bus
}

type asyncTask struct {
	sub *Sub
	msg interface{}
}

func (b *Bus) worker() {
	for {
		select {
		case task := <-b.taskQueue:
			task.sub.receive(task.msg)
			b.taskPool.Put(task) // 将任务放回池中
		case <-b.shutdownSig:
			return
		}
	}
}

func (b *Bus) Subscribe(topic string, sub *Sub) {
	n, _ := b.subNode.LoadOrStore(topic, NewNode())
	node := n.(*node)
	node.rw.Lock()
	defer node.rw.Unlock()
	node.subs = append(node.subs, sub)
}

func (b *Bus) UnSubscribe(topic string, sub *Sub) {
	n, ok := b.subNode.Load(topic)
	if !ok {
		return
	}
	node := n.(*node)
	node.removeSub(sub)
	if node.SubsLen() == 0 {
		b.subNode.Delete(topic)
	}
}

func (b *Bus) Publish(topic string, msg interface{}) error {
	n, ok := b.subNode.Load(topic)
	if !ok {
		return errors.New("topic not exist")
	}
	node := n.(*node)
	node.rw.RLock()
	subs := node.subs // 直接使用原始切片，避免复制
	node.rw.RUnlock()
	if len(subs) == 0 {
		return nil		
	}

	for _, sub := range subs {
		task := b.taskPool.Get().(*asyncTask) // 从池中获取任务
		task.sub = sub
		task.msg = msg
		select {
		case b.taskQueue <- task: // 将任务放入队列
		default:
			// 队列已满，直接清空队列
			// for len(b.taskQueue) > 0 {
			// 	<-b.taskQueue				
			// }
			go sub.receive(msg)
			b.taskPool.Put(task) // 如果任务队列已满，直接执行任务并将任务放回池中
		}
	}
	return nil
}

func (b *Bus) PubFunc(topic string) func(msg interface{}) {
	return func(msg interface{}) {
		b.Publish(topic, msg)
	}
}

func (b *Bus) SubsLen(topic string) (int, error) {
	n, ok := b.subNode.Load(topic)
	if !ok {
		return 0, errors.New("topic not exist")
	}
	node := n.(*node)
	node.rw.RLock()
	defer node.rw.RUnlock()
	return node.SubsLen(), nil
}

type node struct {
	subs []*Sub
	rw   sync.RWMutex
}

func NewNode() *node {
	return &node{
		subs: []*Sub{},
	}
}

func (n *node) SubsLen() int {
	return len(n.subs)
}

func (n *node) removeSub(s *Sub) {
	n.rw.Lock()
	defer n.rw.Unlock()
	idx := n.findSubIdx(s)
	if idx < 0 {
		return
	}
	n.subs = append(n.subs[:idx], n.subs[idx+1:]...)
}

func (n *node) findSubIdx(s *Sub) int {
	for idx, sub := range n.subs {
		if sub == s {
			return idx
		}
	}
	return -1
}

type Sub struct {
	out chan interface{}
}

func NewSub() *Sub {
	return &Sub{
		out: make(chan interface{}, 1024),
	}
}

func (s *Sub) receive(msg interface{}) {
	s.out <- msg
}

func (s *Sub) Out() chan interface{} {
	return s.out
}
